#! /usr/bin/python
# Description
#       - To detect if there are uncorrectable ECC memory errors on GPU cards
#       - If uncorrectable ECC memory errors exist:
#           * reset volatile counter to zero, to make GPU cards working.
#           *
# Author
#       Ran Du <duran@ihep.ac.cn>   v1.0    2020-Apr-21
# Usage
#       <script.py> -h
import subprocess as sp
import sys
import getopt
import re
from datetime import datetime as dt
import json
import logging
import yaml
from os import path
from pprint import pprint


#===============================================================================================
#=============================        global settings       ====================================
#===============================================================================================
# default conf file location, could be reset with cmd options
CONFIG_FILE = "/usr/local/libexec/slurm/detect_reset_ecc_mem.yaml"

# debug switch
#DEBUG = 1 # for dvp
DEBUG = 0 # for production

#===============================================================================================
#=============================         function             ====================================
#===============================================================================================
def usage():
    print("Valid usage is :")
    print("-- ./<script_name> -h")
    print("-- ./<script_name -c <gpu_cards_num> -l <log_file_format> -f <config_file_path>")
    print("     * -c : mandatory")
    print("     * -l : not mandatory, valid value is json")
    print("     * -f : not mandatory, default value is {cfg}".format(cfg=CONFIG_FILE))

def parse_args(args_dict):
    ret = 0
    global config_dict

    # parse cmd args
    cards_n = 0
    log_format = "default"
    config_file = CONFIG_FILE

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hc:l:f:", ["help", "cards=", "logformat=", "cfgfile="])
    except getopt.GetoptError as err:
        sys.stderr.write(err)
        sys.exit(1)

    for opt, arg in opts:
        if opt in ['-h', '--help']:
            usage()
            sys.exit(0)

        elif opt in ['-c', '--cards']:
            assert(re.match("^[0-9]+$",arg)),"ERROR : parse_args : -c, --cards must provide an integer."
            cards_n = int(arg)
            assert(cards_n > 0), "ERROR : parse_args : -c, --cards must be an integer > 0."

        elif opt in ['-l', '--logformat']:
            log_format = arg

        elif opt in ['-f', '--cfgfile']:
            config_file = arg
            assert(path.exists(config_file)), "ERROR : parse_args : {cfg} not existed.".format(cfg=config_file)

        else:
            sys.stderr.write("Invalid option(%s)" % opt)
            usage()
            sys.exit(1)

    # check args
    assert(cards_n > 0), "-c, --cards must be provided."

    # store args
    args_dict['cards_n'] = cards_n
    args_dict['log_format'] = log_format
    args_dict['config_file'] = config_file

    return ret


def detect_uncorr_ecc_err(args_dict, config_dict):
    ret = 0

    # get args
    cards_n = args_dict['cards_n']

    # get output of cmd
    detect_cmd = "nvidia-smi"

    try:
        cmd_out = sp.check_output(detect_cmd, shell=True).decode()
    except sp.CalledProcessError as e:
        cmd_out = e.output
        err_code = e.returncode
        sys.stderr.write(cmd_out)
        return err_code

    # detect if uncorrectable error existed.

    # select lines related with uncorr. ECC memory
    start = 8
    end = start + 4 * cards_n
    cmd_out_lines = cmd_out.split("\n")[start:end]

    card_err_list = []
    card_noerr_list = []
    index = 0
    for index in range(0, len(cmd_out_lines), 4):
        ecc_line = cmd_out_lines[index]
        card_index = ecc_line.split("|")[1].strip().split()[0]
        err_counter = ecc_line.split("|")[3].strip()

        # only when ECC memory enabled and counter > 0
        # uncorrectable ECC memory is detected
        if err_counter.isnumeric():
            err_counter = int(err_counter)
            if err_counter == 0:
                card_noerr_list.append((card_index, err_counter))
            elif err_counter > 0:
                card_err_list.append((card_index, err_counter))

    # with ECC memory errors operations are : log, alarm and reset
    if card_err_list:
        sys.stdout.write("ECC memory errors are deteced.\n\n")

        # log & alarm uncorr. ecc memory errors of the card
        ret = log_alarm_ecc_err(card_err_list, args_dict, config_dict, operation='both')
        if ret != 0:
            sys.stderr.write("ERROR : log_alarm_ecc_err failed.")
            return ret

        # reset error counter
        ret = reset_err_counter(card_err_list, args_dict)
        if ret != 0:
            sys.stderr.write("ERROR : reset_err_counter failed.")
            return ret

    # without ECC memroy errors operations are : log
    if card_noerr_list:
        ret = log_alarm_ecc_err(card_noerr_list, args_dict, config_dict, operation='log')
        if ret != 0:
            sys.stderr.write("ERROR : log_alurm_ecc_err failed.")
            return ret

    return ret

def reset_err_counter(card_err_list, args_dict):
    ret = 0

    # number of gpu cards of the host
    cards_n = int(args_dict['cards_n'])

    for card_err_tuple in card_err_list:
        card_index = card_err_tuple[0]
        assert(re.match("^[0-9]+$", card_index)), \
            "ERROR : reset_err_counter : card_index must be an integer. "
        card_index = int(card_index)

        assert(card_index < cards_n), \
            "ERROR : reset_err_counter : card_index({}) is larger than cards_n({})".\
                format(card_index, cards_n)

        volatile_counter = 0
        reset_cmd = 'nvidia-smi -i {index} --reset-ecc-errors={counter}'.\
            format(index=card_index, counter=volatile_counter)

        ret = sp.call(reset_cmd, shell=True)
        if ret != 0:
            sys.stderr.write("ERROR : reset_err_counter failed")
            return ret

    sys.stdout.write("ECC error counter(s) is/are reset.\n\n")

    return ret

def log_ecc_err_in_json(msg_dict_list, config_dict):
    ret = 0

    log_file = config_dict['log_file']['json']

    with open(log_file, 'a+') as LOG_FILE:
        for msg_dict in msg_dict_list:
            json.dump(msg_dict, LOG_FILE)
            LOG_FILE.write("\n")

    return ret

def log_ecc_err_in_default(msg_dict_list, config_dict):
    ret = 0

    log_file = config_dict['log_file']["default"]

    logging.basicConfig(filename=log_file,
                        level=logging.WARNING,
                        format='%(asctime)s : %(levelname)s : %(message)s')

    for msg_dict in msg_dict_list:
        msg = "{counter} uncorr. ECC memory error(s) on {host} {card} ".\
            format(host=msg_dict['hostname'],
                   card=msg_dict['gpucard'],
                   counter=msg_dict['counter'])

        logging.warning(msg)

    return ret

def log_ecc_err(msg_dict_list, args_dict, config_dict):
    ret = 0

    log_format = args_dict['log_format']
    assert(log_format in config_dict['log_format']), \
        "log_format({fmt}) is not a valid format({fmt_list})".format(fmt=log_format,
                                                                     fmt_list=','.join(config_dict['log_format'])
                                                                     )

    if log_format == 'json':
        ret = log_ecc_err_in_json(msg_dict_list, config_dict)
        if ret != 0:
            sys.stderr.write("ERROR : log_ecc_err_in_json failed.")
            return ret

    if log_format == 'default':
        ret = log_ecc_err_in_default(msg_dict_list, config_dict)
        if ret != 0:
            sys.stderr.write("ERROR : log_ecc_err_logging failed.")
            return ret

    sys.stdout.write("ECC errors are written to log files.\n")

    return ret

def alarm_ecc_err(msg_dict_list, args_dict, config_dict):
    ret = 0

    mail_admin = ' '.join(config_dict['email_admin'])

    mail_title = "Slurm ECC memory error warning"

    mail_context_head = "Dear admin," \
                        "\n\n\tECC memory error(s) are reported on :"

    mail_context_body = ""
    for msg_dict in msg_dict_list:
        time = msg_dict['time']
        host = msg_dict['hostname']
        gpucard = msg_dict['gpucard']
        err_counter = msg_dict['counter']

        msg = "\n\t{time}\t{host}\t{card}\t{counter}".format(time=time, host=host, card=gpucard, counter=err_counter)

        mail_context_body += msg

    mail_context_tail = "\n\tPlease check ASAP." \
                        "\n\nSincerely,\nPomelo"

    mail_context = mail_context_head + mail_context_body + mail_context_tail

    # print only for debug
    if DEBUG:
        print("alarm_ecc_err : mail_context")
        print(mail_context)

    mail_cmd = "echo \"" + mail_context + "\"" + " | mail -s \"" + mail_title + "\" " + mail_admin

    proc = sp.Popen(mail_cmd,
                    shell=True,
                    stdout=sp.PIPE,
                    stderr=sp.PIPE)
    stdout, stderr = proc.communicate()
    ret = proc.returncode
    if ret != 0:
        sys.stderr("ERROR : alarm_ecc_err " + stderr)
        return ret

    sys.stdout.write("Alarm mails are sent.\n")

    return ret


def log_alarm_ecc_err(card_err_list, args_dict, config_dict, operation):
    ret = 0

    try:
        host = sp.check_output('hostname', shell=True).decode().strip()
    except sp.CalledProcessError as e:
        ret = e.returncode
        sys.stderr.write("ERROR : log_alarm_ecc_err : failed to get hostname.")
        return ret

    time = dt.now().strftime("%Y-%m-%d %H:%M:%S")

    msg_dict_list = []
    for card_err_tuple in card_err_list:
        card_index = card_err_tuple[0]
        err_counter = card_err_tuple[1]

        msg_dict = {
            'gpucard': '/dev/nvidia{index}'.format(index=card_index),
            'counter': err_counter,
            'time': time,
            'hostname': host
        }

        msg_dict_list.append(msg_dict)

    assert(operation in ['log','alarm','both']), "ERRROR : log_alarm_ecc_err : Invalid operation."

    if operation in ['log', 'both']:
        # log error counter
        ret = log_ecc_err(msg_dict_list, args_dict, config_dict)
        if ret != 0:
            sys.stderr.write("ERROR : log_ecc_err failed.")
            return ret

    if operation in ['alarm', 'both']:
        # send alarm email to admin
        ret = alarm_ecc_err(msg_dict_list, args_dict, config_dict)
        if ret != 0:
            sys.stderr.write("ERROR : alarm_ecc_err failed.")
            return ret

    return ret

def load_config_files(args_dict):

    config_file = args_dict['config_file']

    with open(config_file, 'r') as CFG:
        config_dict = yaml.load(CFG, Loader=yaml.SafeLoader)

    # print only for debug
    if DEBUG:
        print("load_config_files : config_dict")
        pprint(config_dict)

    sys.stdout.write("Config file are loaded.\n")

    return config_dict

#===============================================================================================
#==================================            main()           ================================
#===============================================================================================
def main():

    # parse cmd args
    args_dict = {}
    ret = parse_args(args_dict)
    if ret != 0:
        sys.stderr.write("ERROR : parse_args failed.")
        sys.exit(ret)

    # load config files
    config_dict = load_config_files(args_dict)
    if config_dict == {}:
        sys.stderr.write("ERROR : load_config_files failed.")
        sys.exit(ret)

    # detect uncorretable ECC memory errors
    ret = detect_uncorr_ecc_err(args_dict, config_dict)
    if ret != 0:
        sys.stderr.write("ERROR : detect_uncorr_ecc_err failed.")
        sys.exit(ret)


if __name__ == '__main__':
    main()

